{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "name": "SENet_CIFAR10.ipynb",
      "provenance": [],
      "collapsed_sections": [],
      "include_colab_link": true
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "accelerator": "GPU"
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "view-in-github",
        "colab_type": "text"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/github/OUCTheoryGroup/colab_demo/blob/master/SENet_CIFAR10.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Vyx9j3WEsCVq",
        "colab_type": "text"
      },
      "source": [
        "## Squeeze-and-Excitation Networks\n",
        "\n",
        "国内自动驾驶创业公司 Momenta 在 ImageNet 2017 挑战赛中夺冠，网络架构为 SENet，论文作者为 Momenta 高级研发工程师胡杰\n",
        "\n",
        "网络是否可以从其他层面来考虑去提升性能，比如考虑特征通道之间的关系？SENet就是考虑了通道层面的关系。具体来说，就是通过学习的方式来获取每个特征通道的重要程度，然后依照这个重要程度去提升有用的特征并抑制对当前任务用处不大的特征。\n",
        "\n",
        "SE模块如下图所示：\n",
        "\n",
        "![替代文字](https://gaopursuit.oss-cn-beijing.aliyuncs.com/202003/20200309171011.jpg)\n",
        "\n",
        "首先是 Squeeze 操作，顺着空间维度来进行特征压缩，将每个二维的特征通道变成一个实数，这个实数某种程度上具有全局的感受野，表征着在特征通道上响应的全局分布\n",
        "\n",
        "然后是 Excitation 操作，类似于循环神经网络门的机制。通过参数 w 来为每个特征通道生成权重，其中参数 w 被学习用来显式地建模特征通道间的相关性。（这里是两个全连接层）\n",
        "\n",
        "最后是 Reweight 的操作，将 Excitation 的输出的权重看做是进过特征选择后的每个特征通道的重要性，然后通过乘法逐通道加权到先前的特征上，完成在通道上对原始特征的重标定。\n",
        "\n",
        "下面是 BasicBlock 的代码："
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "Dhpe8ufGr-Jx",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "import torch\n",
        "import torch.nn as nn\n",
        "import torch.nn.functional as F\n",
        "import torchvision\n",
        "import torchvision.transforms as transforms\n",
        "import matplotlib.pyplot as plt\n",
        "import numpy as np\n",
        "import torch.optim as optim\n",
        "\n",
        "class BasicBlock(nn.Module):\n",
        "    def __init__(self, in_channels, out_channels, stride=1):\n",
        "        super(BasicBlock, self).__init__()\n",
        "        self.conv1 = nn.Conv2d(in_channels, out_channels, kernel_size=3, stride=stride, padding=1, bias=False)\n",
        "        self.bn1 = nn.BatchNorm2d(out_channels)\n",
        "        self.conv2 = nn.Conv2d(out_channels, out_channels, kernel_size=3, stride=1, padding=1, bias=False)\n",
        "        self.bn2 = nn.BatchNorm2d(out_channels)\n",
        "\n",
        "        # shortcut的输出维度和输出不一致时，用1*1的卷积来匹配维度\n",
        "        self.shortcut = nn.Sequential()\n",
        "        if stride != 1 or in_channels != out_channels:\n",
        "            self.shortcut = nn.Sequential(\n",
        "                nn.Conv2d(in_channels, out_channels, kernel_size=1, stride=stride, bias=False),\n",
        "                nn.BatchNorm2d(out_channels))\n",
        "\n",
        "        # 在 excitation 的两个全连接\n",
        "        self.fc1 = nn.Conv2d(out_channels, out_channels//16, kernel_size=1) \n",
        "        self.fc2 = nn.Conv2d(out_channels//16, out_channels, kernel_size=1)\n",
        "\n",
        "    #定义网络结构\n",
        "    def forward(self, x):\n",
        "        #feature map进行两次卷积得到压缩\n",
        "        out = F.relu(self.bn1(self.conv1(x)))\n",
        "        out = self.bn2(self.conv2(out))\n",
        "\n",
        "        # Squeeze 操作：global average pooling\n",
        "        w = F.avg_pool2d(out, out.size(2))\n",
        "        \n",
        "        # Excitation 操作： fc（压缩到16分之一）--Relu--fc（激到之前维度）--Sigmoid（保证输出为 0 至 1 之间）\n",
        "        w = F.relu(self.fc1(w))\n",
        "        w = F.sigmoid(self.fc2(w))\n",
        "\n",
        "        # 重标定操作： 将卷积后的feature map与 w 相乘\n",
        "        out = out * w \n",
        "        # 加上浅层特征图\n",
        "        out += self.shortcut(x)\n",
        "        #R elu激活\n",
        "        out = F.relu(out)\n",
        "        return out"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "loSD-N6ApxpD",
        "colab_type": "text"
      },
      "source": [
        "## SENet 网络"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "9s_SAyhnvEXO",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "#创建SENet网络\n",
        "class SENet(nn.Module):\n",
        "    def __init__(self):\n",
        "        super(SENet, self).__init__()\n",
        "        #最终分类的种类数\n",
        "        self.num_classes = 10\n",
        "        #输入深度为64\n",
        "        self.in_channels = 64\n",
        "\n",
        "        #先使用64*3*3的卷积核\n",
        "        self.conv1 = nn.Conv2d(3, 64, kernel_size=3, stride=1, padding=1, bias=False)\n",
        "        self.bn1 = nn.BatchNorm2d(64)\n",
        "        #卷积层的设置，BasicBlock\n",
        "        #2,2,2,2为每个卷积层需要的block块数\n",
        "        self.layer1 = self._make_layer(BasicBlock,  64, 2, stride=1)\n",
        "        self.layer2 = self._make_layer(BasicBlock, 128, 2, stride=2)\n",
        "        self.layer3 = self._make_layer(BasicBlock, 256, 2, stride=2)\n",
        "        self.layer4 = self._make_layer(BasicBlock, 512, 2, stride=2)\n",
        "        #全连接\n",
        "        self.linear = nn.Linear(512, self.num_classes)\n",
        "\n",
        "    #实现每一层卷积\n",
        "    #blocks为大layer中的残差块数\n",
        "    #定义每一个layer有几个残差块，resnet18是2,2,2,2\n",
        "    def _make_layer(self, block, out_channels, blocks, stride):\n",
        "        strides = [stride] + [1]*(blocks-1)\n",
        "        layers = []\n",
        "        for stride in strides:\n",
        "            layers.append(block(self.in_channels, out_channels, stride))\n",
        "            self.in_channels = out_channels\n",
        "        return nn.Sequential(*layers)\n",
        "\n",
        "    #定义网络结构\n",
        "    def forward(self, x):\n",
        "        out = F.relu(self.bn1(self.conv1(x)))\n",
        "        out = self.layer1(out)\n",
        "        out = self.layer2(out)\n",
        "        out = self.layer3(out)\n",
        "        out = self.layer4(out)\n",
        "        out = F.avg_pool2d(out, 4)\n",
        "        out = out.view(out.size(0), -1)\n",
        "        out = self.linear(out)\n",
        "        return out"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "KAUFxvGYpf2E",
        "colab_type": "text"
      },
      "source": [
        "后面的训练和测试函数和别的里面一样，这里就省去了 ~~~ "
      ]
    }
  ]
}
